/* SPDX-License-Identifier: GPL-2.0 */
static const char *__doc__ = "SRv4 Program Loader\n"
	" - Allows selecting SRv6 section --progsec name to attach to --dev\n";

#include <stdlib.h>
#include <errno.h>
#include <getopt.h>
#include <unistd.h>

#include <bpf/bpf.h>
#include <bpf/libbpf.h>

#include <linux/if_link.h> /* depend on kernel-headers installed */
#include <sys/ioctl.h>
#include <netdb.h>

#include "./common/common_params.h"
#include "./common/common_user_bpf_xdp.h"
#include "./common/common_libbpf.h"
#include "common.h"
#include "ip_mac_conv.h"

static const char *default_filename = "srv4_kern.o";

static const struct option_wrapper long_options[] = {

	{{"help",        no_argument,		NULL, 'h' },
	 "Show help", false},

	{{"dev",         required_argument,	NULL, 'd' },
	 "Operate on device <ifname>", "<ifname>", true},

	{{"skb-mode",    no_argument,		NULL, 'S' },
	 "Install SRv4 program in SKB (AKA generic) mode"},

	{{"native-mode", no_argument,		NULL, 'N' },
	 "Install SRv4 program in native mode"},

	{{"auto-mode",   no_argument,		NULL, 'A' },
	 "Auto-detect SKB or native mode"},

	{{"force",       no_argument,		NULL, 'F' },
	 "Force install, replacing existing program on interface"},

	{{"unload",      no_argument,		NULL, 'U' },
	 "Unload SRv4 program instead of loading"},

	{{"quiet",       no_argument,		NULL, 'q' },
	 "Quiet mode (no output)"},

	{{"filename",    required_argument,	NULL,  1  },
	 "Load program from <file>", "<file>"},

	{{"progsec",    required_argument,	NULL,  2  },
	 "Load program in <section> of the ELF file", "<section>"},

	{{0, 0, NULL,  0 }, NULL, false}
};

#ifndef PATH_MAX
#define PATH_MAX  512	
#endif
#define MAX_MAP_NUM 10

const char *pin_basedir =  "/sys/fs/bpf";
const char *sidmap_name    =  "sidmap";
const char *patmap_name    =  "pathmap";
const char *flwmap_name    =  "flowmap";
const char *macmap_name    =  "macmap";
const char *arpmap_name    =  "arpmap";
const char *rdimap_name    =  "redirect_map";

int 
get_net_mac(__u8 mac[6], const char *ifname)
{
	int sockfd;
	struct ifreq tmp;
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd<0)
	{
		printf("create socket fail\n");
		return -1;
	}
	memset(&tmp, 0, sizeof(struct ifreq));
	memcpy(tmp.ifr_name, ifname, sizeof(tmp.ifr_name)-1);
	if ( (ioctl(sockfd, SIOCGIFHWADDR, &tmp)) < 0 )
	{
		printf("mac ioctl error\n");
		return -1;
	}
	mac[0] = (__u8)tmp.ifr_hwaddr.sa_data[0];
	mac[1] = (__u8)tmp.ifr_hwaddr.sa_data[1];
	mac[2] = (__u8)tmp.ifr_hwaddr.sa_data[2];
	mac[3] = (__u8)tmp.ifr_hwaddr.sa_data[3];
	mac[4] = (__u8)tmp.ifr_hwaddr.sa_data[4];
	mac[5] = (__u8)tmp.ifr_hwaddr.sa_data[5];
	close(sockfd);

	return 0;
}

int
map_lookup(const char *ifname, const char *map_name, struct bpf_map_info *mapinfo)
{
	struct bpf_map_info map_expect = { 0 };
	char pin_dir[PATH_MAX];
	int stats_map_fd;

	int len = snprintf(pin_dir, PATH_MAX, "%s/%s", pin_basedir, ifname);
	if (len < 0) {
		fprintf(stderr, "ERR: creating pin dirname\n");
		return -1;
	}

	map_expect.key_size    = mapinfo->key_size;
	map_expect.value_size  = mapinfo->value_size;
	map_expect.max_entries = mapinfo->max_entries;
	memset(mapinfo, 0, sizeof(struct bpf_map_info));
	stats_map_fd = open_bpf_map_file(pin_dir, map_name, mapinfo);
	if (stats_map_fd < 0) {
		return -1;
	}

	int err = check_map_fd_info(mapinfo, &map_expect);
	if (err) {
		fprintf(stderr, "ERR: map via FD not compatible\n");
		return -1;
	}
	return stats_map_fd;
}

int
get_arp_items(struct ARP arp[MAX_ARP_NUM+1])
{
	unsigned char ipstr[20], type[20], flags[20], macstr[20], mask[20], devstr[20], mac[ETH_ALEN];
	unsigned int ip;
	char *dgw_name="_gateway";
	unsigned int dgw_ip;
	struct hostent *phost = gethostbyname(dgw_name);
	if(!phost)
	{
		puts("Get IP address error!");
		system("pause");
		exit(0);
	}
	char ip4[4];
	memcpy(ip4, phost->h_addr_list[0], 4);
	ip2int(&dgw_ip, ip4);
	dgw_ip = htonl(dgw_ip);

	FILE *fp;
	fp = fopen("/proc/net/arp", "r");
	char line[1000];
	fgets(line, 1000, fp);
	struct ifreq ifreq;
	int sock;
	if((sock=socket(AF_INET,SOCK_STREAM,0))<0)
	{
		perror("socket");
		return 2;
	}
	int k = 1;
	while (fgets(line, 1000, fp) != NULL)
	{
		sscanf(line, "%s %s %s %s %s %s", ipstr, type, flags, macstr, mask, devstr);
		ipstr2int(&ip, ipstr);
		str2mac(mac, macstr);
		arp[k].ifindex = if_nametoindex(devstr);
		memcpy(arp[k].eth_dest, mac, ETH_ALEN);
		arp[k].ip = htonl(ip);

		memcpy(ifreq.ifr_name, devstr, sizeof(devstr));
		if(ioctl(sock,SIOCGIFHWADDR,&ifreq)<0)
		{
			perror("ioctl");
			return 3;
		}
		memcpy(arp[k].eth_source, ifreq.ifr_hwaddr.sa_data, ETH_ALEN);
		if (arp[k].ip == dgw_ip)
		{
			arp[0].ifindex = arp[k].ifindex;
			arp[0].ip = arp[k].ip;
			memcpy(arp[0].eth_dest, arp[k].eth_dest, ETH_ALEN);
			memcpy(arp[0].eth_source, arp[k].eth_source, ETH_ALEN);
		}

		k ++; 
	}
	fclose(fp);

}

static void
update_mac_map(const char *ifname)
{
	//------------- get net mac ... ----------
	__u8 mac[ETH_ALEN];
	if (get_net_mac(mac, ifname))
	{
		printf("get_net_mac of %s error\n", ifname);
		exit(-1);
	}
	//------------- get net mac end ----------

	//------------- find mac map beginning ... ----------
	struct bpf_map_info mapinfo = { 0 };
	mapinfo.key_size = sizeof(__u32);
	mapinfo.value_size = sizeof(struct MAC);
	mapinfo.max_entries = 2;
	int map_fd;
	map_fd = map_lookup(ifname, macmap_name, &mapinfo);
	if (map_fd < 0)
	{
		printf("Do not find map %s\n", macmap_name);
		exit(-1);
	}
	//------------- find mac map end ----------

	//------------- update net mac map beginning ... ----------
	__u32 key = 0;
	struct MAC net_mac;
	memcpy(net_mac.mac, mac, ETH_ALEN);
	if (bpf_map_update_elem(map_fd, &key, &net_mac, 0) < 0) 
	{
		fprintf(stderr, "WARN: Failed to update bpf map file: err(%d):%s\n", errno, strerror(errno));
		exit(-1);
	}
	//------------- update mac map end ----------
	
}

static void
update_arp_map(const char *ifname)
{
	//------------- get arp... ----------
	struct ARP arp_items[MAX_ARP_NUM+1];
	memset(arp_items, 0, sizeof(struct ARP)*(MAX_ARP_NUM+1));
	get_arp_items(arp_items);
	//------------- get arp end ----------

	//------------- find arp map beginning ... ----------
	struct bpf_map_info mapinfo = { 0 };
	mapinfo.key_size = sizeof(__u32);
	mapinfo.value_size = sizeof(struct ARP);
	mapinfo.max_entries = MAX_ARP_NUM;
	int map_fd;
	map_fd = map_lookup(ifname, arpmap_name, &mapinfo);
	if (map_fd < 0)
	{
		printf("Do not find map %s\n", arpmap_name);
		exit(-1);
	}
	//------------- find arp map end ----------

	//------------- update arp map beginning ... ----------
	struct ARP arp;
	__u32 key;
	arp.ifindex = arp_items[0].ifindex;
	arp.ip = arp_items[0].ip;
	memcpy(arp.eth_source, arp_items[0].eth_source, ETH_ALEN);
	memcpy(arp.eth_dest, arp_items[0].eth_dest, ETH_ALEN);
	key = 0;
	if (bpf_map_update_elem(map_fd, &key, &arp, BPF_ANY) < 0) 
	{
		fprintf(stderr, "WARN: Failed to update bpf map file: err(%d):%s\n", errno, strerror(errno));
		exit(-1);
	}
	for (int i=0; i<MAX_ARP_NUM+1; i++)
	{
		if (arp_items[i].ip == 0)
			continue;
		arp.ifindex = arp_items[i].ifindex;
		arp.ip = arp_items[i].ip;
		memcpy(arp.eth_source, arp_items[i].eth_source, ETH_ALEN);
		memcpy(arp.eth_dest, arp_items[i].eth_dest, ETH_ALEN);
		key = arp_items[i].ip;
		if (bpf_map_update_elem(map_fd, &key, &arp, 0) < 0) 
		{
			fprintf(stderr, "WARN: Failed to update bpf map file: err(%d):%s\n", errno, strerror(errno));
			exit(-1);
		}
	}
	//------------- update arp map end ----------
	
}

static void
update_redirect_map(const char *ifname)
{
	//------------- get arp... ----------
	struct ARP arp_items[MAX_ARP_NUM+1];
	memset(arp_items, 0, sizeof(struct ARP)*(MAX_ARP_NUM+1));
	get_arp_items(arp_items);
	//------------- get arp end ----------

	//------------- find redirect map beginning ... ----------
	struct bpf_map_info mapinfo = { 0 };
	mapinfo.key_size = sizeof(__u32);
	mapinfo.value_size = sizeof(struct REDIRECT);
	mapinfo.max_entries = MAX_ARP_NUM;
	int map_fd;
	map_fd = map_lookup(ifname, rdimap_name, &mapinfo);
	if (map_fd < 0)
	{
		printf("Do not find map %s\n", rdimap_name);
		exit(-1);
	}
	//------------- find redirect map end ----------

	//------------- update redirect map beginning ... ----------
	struct REDIRECT rdi;
	__u32 key;
	rdi.valid = 1;
	rdi.egr_ifi = arp_items[0].ifindex;
	rdi.ipv4 = arp_items[0].ip;
	memcpy(rdi.eth_source, arp_items[0].eth_source, ETH_ALEN);
	memcpy(rdi.eth_dest, arp_items[0].eth_dest, ETH_ALEN);
	key = 0;
	if (bpf_map_update_elem(map_fd, &key, &rdi, 0) < 0) 
	{
		fprintf(stderr, "WARN: Failed to update bpf map file: err(%d):%s\n", errno, strerror(errno));
		exit(-1);
	}
	for (int i=0; i<MAX_ARP_NUM+1; i++)
	{
		if (arp_items[i].ip == 0)
			continue;

		rdi.valid = 1;
		rdi.egr_ifi = arp_items[i].ifindex;
		rdi.ipv4 = arp_items[i].ip;
		memcpy(rdi.eth_source, arp_items[i].eth_source, ETH_ALEN);
		memcpy(rdi.eth_dest, arp_items[i].eth_dest, ETH_ALEN);

		key = i + 1;
		if (bpf_map_update_elem(map_fd, &key, &rdi, 0) < 0) 
		{
			fprintf(stderr, "WARN: Failed to update bpf map file: err(%d):%s\n", errno, strerror(errno));
			exit(-1);
		}
	}
	//------------- update arp map end ----------

}


/* Pinning maps under /sys/fs/bpf in subdir */
int pin_maps_in_bpf_object(struct bpf_object *bpf_obj, const char *subdir, const char *map_name)
{
	char map_filename[PATH_MAX];
	char pin_dir[PATH_MAX];
	int err, len;

	len = snprintf(pin_dir, PATH_MAX, "%s/%s", pin_basedir, subdir);
	if (len < 0) {
		fprintf(stderr, "ERR: creating pin dirname\n");
		return EXIT_FAIL_OPTION;
	}

	len = snprintf(map_filename, PATH_MAX, "%s/%s/%s",
		       pin_basedir, subdir, map_name);
	if (len < 0) {
		fprintf(stderr, "ERR: creating map_name\n");
		return EXIT_FAIL_OPTION;
	}

	if (access(map_filename, F_OK ) != -1 ) {
		if (verbose)
			printf(" - Unpinning (remove) prev maps in %s/\n",
			       pin_dir);

		err = bpf_object__unpin_maps(bpf_obj, pin_dir);
		if (err) {
			fprintf(stderr, "ERR: UNpinning maps in %s\n", pin_dir);
			return EXIT_FAIL_BPF;
		}
	}
	if (verbose)
		printf(" - Pinning maps in %s/\n", pin_dir);

	err = bpf_object__pin_maps(bpf_obj, pin_dir);
	if (err)
		return EXIT_FAIL_BPF;

	return 0;
}

int main(int argc, char **argv)
{
	struct bpf_object *bpf_obj;
	int err;
		char pin_dir[PATH_MAX];
		char cmd[PATH_MAX+10];

	struct config cfg = {
		.xdp_flags = XDP_FLAGS_UPDATE_IF_NOEXIST | XDP_FLAGS_DRV_MODE,
		.ifindex   = -1,
		.do_unload = false,
	};
	strncpy(cfg.filename, default_filename, sizeof(cfg.filename));
	parse_cmdline_args(argc, argv, long_options, &cfg, __doc__);

	/* Required option */
	if (cfg.ifindex == -1) {
		fprintf(stderr, "ERR: required option --dev missing\n\n");
		usage(argv[0], __doc__, long_options, (argc == 1));
		return EXIT_FAIL_OPTION;
	}
	if (cfg.do_unload) {

		xdp_link_detach(cfg.ifindex, cfg.xdp_flags, 0);
		int len = snprintf(pin_dir, PATH_MAX, "%s/%s", pin_basedir, cfg.ifname);
		sprintf(cmd, "rm -rf %s\n", pin_dir);
		system(cmd);
		return 0;
	}

	bpf_obj = load_bpf_and_xdp_attach(&cfg);
	if (!bpf_obj)
		return EXIT_FAIL_BPF;


		printf("Success: Loaded BPF-object(%s) and used section(%s)\n",
		       cfg.filename, cfg.progsec);
		printf(" - XDP prog attached on device:%s(ifindex:%d)\n",
		       cfg.ifname, cfg.ifindex);
#if 1
	//-------------- pin maps to file  begin...--------------------------
	err = pin_maps_in_bpf_object(bpf_obj, cfg.ifname, sidmap_name);
	if (err) {
		fprintf(stderr, "ERR: pinning %s\n", sidmap_name);
		return err;
	}
	err = pin_maps_in_bpf_object(bpf_obj, cfg.ifname, patmap_name);
	if (err) {
		fprintf(stderr, "ERR: pinning %s\n", patmap_name);
		return err;
	}
	err = pin_maps_in_bpf_object(bpf_obj, cfg.ifname, flwmap_name);
	if (err) {
		fprintf(stderr, "ERR: pinning %s\n", flwmap_name);
		return err;
	}
	err = pin_maps_in_bpf_object(bpf_obj, cfg.ifname, arpmap_name);
	if (err) {
		fprintf(stderr, "ERR: pinning %s\n", arpmap_name);
		return err;
	}
	update_mac_map(cfg.ifname);
	update_arp_map(cfg.ifname);
	update_redirect_map(cfg.ifname);
	//-------------- pin maps to file  end --------------------------
#endif

	return EXIT_OK;
}
